"""
@Filename       : lastfm_preprocess.py
@Create Time    : 2022/3/5 15:59
@Author         : Rylynn
@Description    : 

"""
import datetime
import random

cascade_dict = {}
cascade_user_time = {}
user_count = {}
content_dict = {}

with open('../../../data/raw/userid-timestamp-artid-artname-traid-traname.tsv', 'r+', encoding='utf8') as f:
    for line in f.readlines():
        user_id, timestamp, artist_id, artist, track_id, track = line.strip().split('\t')
        timestamp = int(datetime.datetime.strptime(timestamp, '%Y-%m-%dT%H:%M:%SZ').timestamp())
        content_dict[track_id] = artist + ' ' + track
        if track_id == '' or user_id == '':
            continue

        if cascade_dict.get(track_id):
            cascade_dict[track_id].append((user_id, timestamp))
        else:
            cascade_dict[track_id] = [(user_id, timestamp)]

        if user_count.get(user_id):
            user_count[user_id] += 1
        else:
            user_count[user_id] = 1

print(len(cascade_dict))

final_user_set = set()
for user, count in user_count.items():
    if count > 5:
        final_user_set.add(user)


new_cascade_dict = {}

for id, cascade in cascade_dict.items():
    new_cascade = []

    for u, t in cascade:
        if u in final_user_set:
            new_cascade.append((u, t))

    sorted_cascade = sorted(new_cascade, key=lambda x: x[1])

    removed_redundant_cascade = []
    existing_user = set()
    for u, t in cascade:
        if not u in existing_user:
            removed_redundant_cascade.append((u, t))
            existing_user.add(u)

    if len(removed_redundant_cascade) > 300 or len(removed_redundant_cascade) < 2:
        continue

    new_cascade_dict[id] = removed_redundant_cascade

random.seed(2)

print(len(new_cascade_dict.values()))
print(sum(map(lambda c: len(c[1]), new_cascade_dict.items()))/len(new_cascade_dict.values()))

sampled_cascade_dict = random.sample(new_cascade_dict.items(), 10000)

with open('../../../data/lastfm/content.txt', mode='w', encoding='utf-8') as f:
    for track_id, c in sampled_cascade_dict:
        f.write('{}\t{}\n'.format(track_id, content_dict[track_id]))

def write_cascade(cascade, filename):
    with open(filename, 'w') as f:
        for k, c in cascade:
            f.write('{} '.format(k))
            for u, t in c:
                f.write('{},{} '.format(u, t))
            f.write('\n')

write_cascade(sampled_cascade_dict[:8000], '../../../data/lastfm/cascade.txt')
write_cascade(sampled_cascade_dict[8000:9000], '../../../data/lastfm/cascadetest.txt')
write_cascade(sampled_cascade_dict[9000:], '../../../data/lastfm/cascadevalid.txt')


